import os
import torch
from torch.autograd import Variable
from torchvision import transforms
import torch.nn.functional as F

import numpy as np
import cv2
from model import get_net

os.environ['CUDA_VISIBLE_DEVICES'] = '0'


def main():

    model_name = 'u2net_cls'
    image_dir = './test_data/gt_test2.txt'
    prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep)
    model_dir = os.path.join(os.getcwd(), 'saved_models', model_name + os.sep + 'u2net_cls_epoch719_bce_itr_163000_train_4.812116.pth')

    # net = get_net(model_name)
    #
    # if torch.cuda.is_available():
    #     net.load_state_dict(torch.load(model_dir))
    #     net.cuda()
    # else:
    #     net.load_state_dict(torch.load(model_dir, map_location='cpu'))
    net = torch.jit.load('C:/Users/WDSMY/Desktop/removal/classifier2/data/classifier_model.pt', map_location=torch.device('cuda:0'))
    net.eval()

    inputs_test = cv2.imread('./test_data/gt_cls2/wmtext_2nd_3.png')
    inputs_test = transforms.ToTensor()(inputs_test)
    c, h, w = inputs_test.shape
    inputs_test = inputs_test.view(1, c, h, w)


    if torch.cuda.is_available():
        inputs_test = Variable(inputs_test.cuda())
    else:
        inputs_test = Variable(inputs_test)

    d0, d1, d2, d3, d4, d5, d6, dc = net(inputs_test)


    pred_c = F.softmax(dc, dim=1)[:,1].tolist()
    _predict = np.array(pred_c)
    _predict = np.where(_predict>0.5, 1, 0)
    print(_predict[0])


if __name__ == "__main__":
    main()

